package com.zyx.flinkdemo.stream.sink.hbasebatch;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.zyx.flinkdemo.pojo.HttpDataModel;
import org.apache.commons.cli.*;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.AllWindowFunction;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.flink.util.Collector;
import org.apache.kafka.clients.consumer.ConsumerConfig;

import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

/**
 * @author Yaxi.Zhang
 * @since 2021/8/10 13:10
 * desc: 读取Kafka写入HBase主类
 */
public class FlinkToHbaseApp {
    public static void main(String[] args) throws Exception {
        // kafka 需要的参数
        String brokers = "127.0.0.1:9092";
        String topic = "test";
        // hbase 需要的参数
        String hbaseZookeeperHost = "hbase";
        String hbaseZookeeperPort = "2181";

        // 接收命令行参数，覆盖默认值
        Options options = new Options();
        options.addOption("kafka_brokers", true, "kafka cluster hosts, such 127.0.0.1:9092");
        options.addOption("kafka_topic", true, "kafka cluster topic, default: test");
        options.addOption("hbase_zookeeper_host", true, "hbase zookeeper host, default: hbase");
        options.addOption("hbase_zookeeper_port", true, "hbase zookeeper port, default: 2181");

        CommandLineParser parser = new DefaultParser();
        CommandLine line = parser.parse(options, args);

        if (line.hasOption("kafka_brokers")) {
            brokers = line.getOptionValue("kafka_brokers");
        } else {
            HelpFormatter formatter = new HelpFormatter();
            formatter.printHelp("flink write hbase job", options);
            System.exit(1);
        }
        if (line.hasOption("kafka_topic")) {
            topic = line.getOptionValue("kafka_topic");
        }
        if (line.hasOption("hbase_zookeeper_host")) {
            hbaseZookeeperHost = line.getOptionValue("hbase_zookeeper_host");
        }
        if (line.hasOption("hbase_zookeeper_port")) {
            hbaseZookeeperPort = line.getOptionValue("hbase_zookeeper_port");
        }

        // 执行任务
        doExcute(brokers, topic, hbaseZookeeperHost, hbaseZookeeperPort);
    }

    /**
     * 具体任务执行
     */
    public static void doExcute(String kafkaBrokers, String topic, String hbaseZookeeperHost, String hbaseZookeeperPort) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // 设置 kafka source
        env.enableCheckpointing(5000 * 100000);

        Properties props = getKafkaProperties();
        props.setProperty("bootstrap.servers", kafkaBrokers);
        DataStream<String> stream = env.addSource(new FlinkKafkaConsumer<String>(topic, new SimpleStringSchema(), props));

        // 过滤不标准格式的数据，并格式化
        DataStream<HttpDataModel> formatedStream = stream.filter(s -> {
            JSONObject obj = JSONObject.parseObject(s);
            return obj.containsKey("project") && obj.containsKey("table") && obj.containsKey("data");
        }).map(s -> JSON.parseObject(s, HttpDataModel.class));

        // 在 10 秒的时间窗口内，每 100 条触发输出到 hbase
        SingleOutputStreamOperator<List<HttpDataModel>> batchStream = formatedStream
                .timeWindowAll(Time.seconds(10))
                .trigger(new CountTrigger(100))
                .apply(new AllWindowFunction<HttpDataModel, List<HttpDataModel>, Window>() {
                    @Override
                    public void apply(Window window, Iterable<HttpDataModel> values, Collector<List<HttpDataModel>> out) throws Exception {
                        List<HttpDataModel> lists = new ArrayList<HttpDataModel>();
                        for (HttpDataModel value : values) {
                            lists.add(value);
                        }
                        out.collect(lists);
                    }
                });

        batchStream.addSink(new HbaseSink(hbaseZookeeperHost, hbaseZookeeperPort));

        // 控制台输出
        /*batchStream.print();*/

        env.execute("integration-http");
    }

    /**
     * 获取 kafka 的默认配置
     */
    public static Properties getKafkaProperties() {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "127.0.0.1:9092");
        props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "dataworks-integration");
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, "1000");
        props.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "30000");
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");

        return props;
    }
}
